import numpy
import pycuda.driver as drv

from spa import MatchTextureModule
from spa import ManyRotate2D

# Projection matching reference texture
class PMRefTexture:
	def __init__(self,rimages,nrot,radius,nsample):
		self.fill_rotations(rimages,nrot,radius,nsample)
		self.set_device_array()
		self.set_handle()
		self.set_texture()

	def set_device_array(self):
		descriptor = drv.ArrayDescriptor()
		descriptor.height = self.host_array.shape[0]
		descriptor.width = self.host_array.shape[1]
		descriptor.format = drv.dtype_to_array_format(self.host_array.dtype)
		descriptor.num_channels = 1
		self.device_array = drv.Array(descriptor)

	def set_handle(self):
		self.handle = drv.Memcpy2D()
		self.handle.set_src_host(self.host_array)
		self.handle.set_dst_array(self.device_array)
		self.handle.width_in_bytes = self.host_array.strides[0]
		self.handle.src_pitch = self.host_array.strides[0]
		self.handle.dst_pitch = self.host_array.strides[0]
		self.handle.height = self.host_array.shape[0]

	def set_texture(self):
		self.ref_texture = MatchTextureModule.get_texref("ref_texture")
		self.ref_texture.set_array(self.device_array)
		self.ref_texture.set_address_mode(0, drv.address_mode.CLAMP)
		self.ref_texture.set_address_mode(1, drv.address_mode.CLAMP)
		self.ref_texture.set_filter_mode(drv.filter_mode.LINEAR)


	def fill_rotations(self,refs,nrot,radius,nsamp):
		P,N,x = refs.shape
		self.host_array = numpy.zeros((P*N,nrot*N),numpy.float32)

		image = numpy.zeros((N,N),numpy.float32)
		image_gpu = drv.mem_alloc(image.nbytes)

		rimages = numpy.zeros((nrot,N,N),numpy.float32)
		rimages_gpu = drv.mem_alloc(rimages.nbytes)

		angle = 2.0*numpy.pi/float(nrot)

		for i in range(P):
			image = refs[i].copy().astype(numpy.float32)
			drv.memcpy_htod(image_gpu,image)

			rimages.fill(0.0)
			drv.memcpy_htod(rimages_gpu,rimages)

			ManyRotate2D(image_gpu, rimages_gpu, numpy.float32(angle),
				numpy.int32(N), numpy.int32(nsamp), numpy.float32(radius),
				block=(N,1,1), grid=(N,nrot))

			drv.memcpy_dtoh(rimages, rimages_gpu)

			for j in range(nrot):
				self.host_array[i*N:(i+1)*N,j*N:(j+1)*N] = rimages[j]


class PMImgTexture:
	def __init__(self,A):
		self.fill_images(A)
		self.set_device_array()
		self.set_handle()
		self.set_texture()


	def fill_images(self,imgs):
		P,N,x = imgs.shape
		self.host_array = numpy.zeros((N,P*N),numpy.float32)
		for j in range(P):
			self.host_array[:,j*N:(j+1)*N] = imgs[j]

	def set_device_array(self):
		descriptor = drv.ArrayDescriptor()
		descriptor.height = self.host_array.shape[0]
		descriptor.width = self.host_array.shape[1]
		descriptor.format = drv.dtype_to_array_format(self.host_array.dtype)
		descriptor.num_channels = 1
		self.device_array = drv.Array(descriptor)

	def set_handle(self):
		self.handle = drv.Memcpy2D()
		self.handle.set_src_host(self.host_array)
		self.handle.set_dst_array(self.device_array)
		self.handle.width_in_bytes = self.host_array.strides[0]
		self.handle.src_pitch = self.host_array.strides[0]
		self.handle.dst_pitch = self.host_array.strides[0]
		self.handle.height = self.host_array.shape[0]

	def set_texture(self):
		self.img_texture = MatchTextureModule.get_texref("img_texture")
		self.img_texture.set_array(self.device_array)
		self.img_texture.set_address_mode(0, drv.address_mode.CLAMP)
		self.img_texture.set_address_mode(1, drv.address_mode.CLAMP)
		self.img_texture.set_filter_mode(drv.filter_mode.POINT)


class PMWtTexture:
	def __init__(self,A):
		descriptor = drv.ArrayDescriptor()
		descriptor.height = A.shape[0]
		descriptor.width = A.shape[1]
		descriptor.format = drv.dtype_to_array_format(A.dtype)
		descriptor.num_channels = 1
		self.device_array = drv.Array(descriptor)

		self.handle = drv.Memcpy2D()
		self.handle.set_src_host(A)
		self.handle.set_dst_array(self.device_array)
		self.handle.width_in_bytes = A.strides[0]
		self.handle.src_pitch = A.strides[0]
		self.handle.dst_pitch = A.strides[0]
		self.handle.height = A.shape[0]

		self.wt_texture = MatchTextureModule.get_texref("wt_texture")
		self.wt_texture.set_array(self.device_array)
		self.wt_texture.set_address_mode(0, drv.address_mode.CLAMP)
		self.wt_texture.set_address_mode(1, drv.address_mode.CLAMP)
		self.wt_texture.set_filter_mode(drv.filter_mode.POINT)
